# -*- coding: utf-8 -*-
# @Time : 2024-11-2024/11/30 16:52
# @Author : 林枫
# @File : detect.py

import os
import cv2
import numpy as np
from ultralytics import YOLO

# 导入自定义模块
from utils.angle_utils import estimate_3d_from_2d, calculate_angle_3d, get_joint_angles_config, get_skeleton_connections
from utils.visualization import draw_skeleton, add_frame_info


def workouts(model_path, video_path, show=True, show_keypoint_idx=True):
    """
    主处理函数：处理视频，检测姿态，计算角度并可视化结果
    
    参数:
    - model_path: YOLO模型路径
    - video_path: 输入视频路径
    - show: 是否显示处理结果
    - show_keypoint_idx: 是否显示关键点索引
    """
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    assert cap.isOpened(), "Error reading video file"
    
    # 获取视频属性
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
    
    # 设置输出路径
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    out_dir = "./runs"
    os.makedirs(out_dir, exist_ok=True)
    
    # 创建视频写入器
    out_path = f"{out_dir}/{video_name}.avi"
    skeleton_path = f"{out_dir}/{video_name}_skeleton.avi"
    video_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
    skeleton_video_writer = cv2.VideoWriter(skeleton_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))

    # 初始化YOLO姿态检测模型
    model = YOLO(model_path)
    
    # 获取关节角度配置和骨架连接关系
    joint_angles = get_joint_angles_config()
    connections = get_skeleton_connections()

    # 开始逐帧处理视频
    frame_count = 0
    while cap.isOpened():
        # 读取一帧
        success, im0 = cap.read()
        if not success:
            print("视频帧为空或视频处理已成功完成。")
            break
        
        # 更新帧计数
        frame_count += 1
        if frame_count % 30 == 0:  # 每30帧打印一次信息
            print(f"处理第 {frame_count} 帧")
        
        # 创建骨架图像 (黑色背景)
        skeleton_img = np.zeros((h, w, 3), dtype=np.uint8)
        
        # 使用YOLO模型进行姿态检测
        results = model(im0, verbose=False)
        
        # 处理检测结果
        if results and len(results) > 0 and len(results[0]) > 0:
            for result in results[0]:
                try:
                    # 获取关键点数据
                    keypoints = result.keypoints
                    if keypoints is not None and hasattr(keypoints, 'xy') and keypoints.xy is not None and len(keypoints.xy) > 0:
                        # 获取2D关键点
                        keypoints_2d = keypoints.xy.cpu().numpy()
                        
                        try:
                            # 从2D估计3D关键点
                            estimated_3d = estimate_3d_from_2d(keypoints_2d)
                            
                            # 存储计算出的角度
                            angles = {}
                            
                            # 计算各个关节角度
                            for angle_info in joint_angles:
                                p1_idx, p2_idx, p3_idx = angle_info["points"]
                                name = angle_info["name"]
                                
                                # 确保关键点索引有效
                                if (p1_idx < estimated_3d.shape[0] and 
                                    p2_idx < estimated_3d.shape[0] and 
                                    p3_idx < estimated_3d.shape[0]):
                                    
                                    p1 = estimated_3d[p1_idx]
                                    p2 = estimated_3d[p2_idx]
                                    p3 = estimated_3d[p3_idx]
                                    
                                    angle = calculate_angle_3d(p1, p2, p3)
                                    angles[name] = angle
                            
                            # 绘制带有骨架和角度信息的图像
                            if keypoints_2d.shape[0] > 0:
                                keypoints_flat = keypoints_2d[0]  # 取第一个人的关键点
                                
                                # 绘制骨架图像
                                skeleton_img = draw_skeleton(
                                    skeleton_img, 
                                    keypoints_flat, 
                                    connections, 
                                    joint_angles,
                                    angles,
                                    thickness=2,
                                    show_keypoint_idx=show_keypoint_idx
                                )
                                
                                # 绘制原图中的重要角度
                                # 使用YOLO的绘图功能显示关键点
                                im0 = result.plot()
                                
                                # 添加重要角度信息到原图
                                important_angles = ["右膝", "左膝", "躯干前倾", "右肘", "左肘"]
                                for name in important_angles:
                                    if name in angles:
                                        angle_info = next((item for item in joint_angles if item["name"] == name), None)
                                        if angle_info and keypoints_2d.shape[1] > angle_info["points"][1]:
                                            p2_idx = angle_info["points"][1]
                                            pos = tuple(map(int, keypoints_2d[0, p2_idx]))
                                            text = f"{name}: {angles[name]:.1f}°"
                                            
                                            # 添加文本 (避免与检测框重叠)
                                            y_offset = pos[1] - 25
                                            from utils.image_utils import cv2_add_chinese_text
                                            im0 = cv2_add_chinese_text(
                                                im0, text, (pos[0], y_offset), 
                                                font_size=18, color=angle_info["color"]
                                            )
                                
                        except Exception as e:
                            print(f"角度计算或绘制出错: {e}")
                        
                except AttributeError as e:
                    print(f"无法获取关键点: {e}")
        
        # 添加帧信息
        im0 = add_frame_info(im0, frame_count)
        skeleton_img = add_frame_info(skeleton_img, frame_count)
        
        # 显示结果
        if show:
            cv2.imshow("Workout Analysis", im0)
            cv2.imshow("Skeleton Only", skeleton_img)
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord('s'):  # 按's'键保存当前帧
                save_dir = f"{out_dir}/frames"
                os.makedirs(save_dir, exist_ok=True)
                cv2.imwrite(f"{save_dir}/frame_{frame_count}.jpg", im0)
                cv2.imwrite(f"{save_dir}/frame_{frame_count}_skeleton.jpg", skeleton_img)
                print(f"保存了第 {frame_count} 帧图像")
        
        # 保存视频
        video_writer.write(im0)
        skeleton_video_writer.write(skeleton_img)

    # 清理资源
    cv2.destroyAllWindows()
    video_writer.release()
    skeleton_video_writer.release()
    cap.release()


if __name__ == '__main__':
    model_path = "./weights/yolo11x-pose.pt"
    video_path = "./video/跑步.mp4"
    workouts(model_path, video_path)

    # 可以通过取消以下注释来处理其他视频
    # video_path = "./video/俯卧撑.mp4"
    # workouts(model_path, video_path)
    
    # video_path = "./video/引体向上1.mp4"
    # workouts(model_path, video_path)
